;;;; -*- Mode: Lisp; Package: :lib.gl; -*-

(in-package :lib.gl)

;;; Might be better to move a lot of this math stuff into its own package...

;; Defining the variables:
;; -----------------------
;;
;; Definition of an LP (mix of [Ebe03] and http://en.wikipedia.org/wiki/Linear_programming):
;; Ax <= b - Constraints
;; f(x) = transpose(c) x - Function being maximized.
;; u = Vector of slack variables for the dual problem
;; s = Vector of slack variables for the primal problem (the f(x) = cTx and Ax <= b)
;; 
;; Statement of the primal problem:
;;   Maximize f(x) = transpose(c)*x subject to x>=0 and A*x<=b
;; Statement of the dual problem:
;;   Minimize g(y) = transpose(b)*y subject to y>=0 and transpose(A)*y>=c
;;
;; x is feasible if it satisfies the inequality constraints, i.e. for which Ax<=b.  Similar for feasible ys.
;;     | 0  | transpose(A) |
;; M = |----|--------------| S in upper right for quadratic programming.
;;     | -A |      0       |
;;
;;     | -c |
;; q = |----|
;;     |  b |
;;
;;     | u |
;; w = |---| (These variables get renamed to w_i)
;;     | s |
;;
;;     | x |
;; z = | - | (These variables get renamed to z_i)
;;     | y |
;;
;;
;; LCP is defined as
;; w - Mz = q
;; w >> 0, z >> 0 (where x >> y means x_j - y_j is nonnegative for all j)
;; and w_i*z_i = 0 for all i.
;;

(defstruct tableu
  matrix
  decision-var-cnt
  constrain-cnt)

;; tableu is [ A | I | b ]
(defun simplex-pivot (tableu r s)
  (let ((a_rs (mref tableu r s)))
    ;; Step 1
    ;; a_rj <-- a_rj/a_rs for j=[1,n+m]; i.e. gaussian row scale by 1/a_rj
    ;; n = Number of decision variables (i.e. size of vector x)
    ;; m = Number of constraints
    (setf tableu (gaussian-row-scale tableu r (/ 1 a_rs)))
    ;; Step 2 & 3
    ;; a_rs <-- 1 (achieved in 1!)
    ;; a_is <-- 0; i.e. gaussian row elimination
    ;; a_ij <-- a_ij - a_rj * a_is / a_rs
    (iter (for i from 0 below (matrix-rows tableu))
	  (unless (= i r)
	    (setf tableu (gaussian-row-scale-sum tableu r (- (mref tableu i s)) i))))
    tableu))

(defun swap-columns (tableu leaving entering)
  "Swap column named by <leaving> (integer index) with column <entering>.  The parameter
names are chosen to be evocative of Simplex method terminology."
  (let ((ecol (get-column tableu entering)))
    (set-column tableu entering (get-column tableu leaving))
    (set-column tableu leaving ecol)))

;;; Think of the above more as functions for
;;; debugging and aids for hand-calculation.
;;; ----------------------------------------
;;; Now begins the real work of LP!

(defun form-basis-partition (a nconstraints nvars)
  (let ((b (make-matrix nconstraints nconstraints))
	(n (make-matrix nconstraints (- nvars nconstraints))))
    (dotimes (i nvars (values b n))
      (if (< i nconstraints)
	  (set-column b i (get-column a i))
	  (set-column n (- i nconstraints) (get-column a i))))))
	   
(defun augmented-constraints-matrix (a slack-var-cnt)
  (let* ((m (matrix-rows a))
	 (n (matrix-cols a))
	 (new-a (make-matrix m (+ slack-var-cnt n)))
	 (ident (mneg (identity-matrix slack-var-cnt))))
    (iter (for i from 0 below (matrix-cols new-a))
	  (if (< i n)
	      (set-column new-a i (get-column a i))
	      (set-column new-a i (get-column ident (- i n)))))
    new-a))

(defun augmented-cost-vector (c slack-var-cnt)
  (let* ((n (matrix-cols c))
	 (naug (+ n slack-var-cnt))
	 (caug (make-matrix 1 naug)))
    (iter (for i from 0 below naug)
	  (if (< i n)
	      (setf (mref caug 0 i) (mref c 0 i))
	      (setf (mref caug 0 i) 0.0d0)))
    caug))

(defun split-cost-vector (c m n)
  (let ((cb (make-matrix 1 m))
	(cn (make-matrix 1 n)))
    (dotimes (i (+ m n) (values cb cn))
      (if (< i m)
	  (setf (mref cb 0 i) (mref c 0 i))
	  (setf (mref cn 0 (- i m)) (mref c 0 i))))))

(defun find-entering-variable (cost-vector)
  (iter (for x in-vector cost-vector)
	(for i from 0)
	(finding i minimizing x)))

(defun find-leaving-variable (xb v)
  (iter (for num in-vector xb)
	(for denom in-vector v)
	(for i from 0)
	(if (plusp denom)
	    (finding i minimizing (/ num denom)))))

;;
;; This is the 'revised simplex method' described in [Strang 1988].
;; He goes on to state that 'many simplex codes use the "product form of the inverse"'
;; which avoids updating B^-1 on every iteration.  Another method would involve
;; calculating LU factorization for B and solve the three equations:
;;  labmda*B=cb B*v=u B*xb =b
;; This is a "newer" approach (i.e. new for 1988).
;;

(defun naive-simplex-solver (A c b nx ns)
  "Naive LP solver.  Naive because we calculate B^-1 each iteration, which can be
expensive if rank(B) > 3."
  (assert (= (matrix-cols A) nx) (A nx) "Matrix ~S should have ~S columns.")
  (let ((m (matrix-rows A)))
    (multiple-value-bind (basic nonbasic) (form-basis-partition (augmented-constraints-matrix A ns)
						     m (+ nx ns))
      (multiple-value-bind (cb cn) (split-cost-vector (augmented-cost-vector c ns) ns nx)
	(iter (for binv = (inverse basic))
	      (for gamma = (m* cb binv))
	      (for r = (m- cn (m* gamma nonbasic)))
	      (for xb = (m* binv b))
	      (finding xb such-that (vector-plus-p r))
	      (after-each
	       (let* ((ri (find-entering-variable r))
		      (u (get-column nonbasic ri))
		      (v (m* binv u))
		      (k (find-leaving-variable xb v))
		      (cbk (mref cb 0 k)))
		 (if (null k) (leave 'unbounded))
		 ;; The Pivot
		 (set-column nonbasic ri (get-column basic k))
		 (setf (mref cb 0 k) (mref cn 0 ri))
		 (set-column basic k u)
		 (setf (mref cn 0 ri) cbk))))))))

;;;; --- Solving a Linear Complemenary Problem using Lemke's Method ---

(define-condition ray-termination (error)
  ())

(defun lexico-feasible (b)
  (iter (for i from 0 below (matrix-rows b))
	(always (lexico-positive-p (get-row b i)))))

(defun build-tableu (m q)
  (let* ((zcnt (matrix-cols m))
	 (w (identity-matrix zcnt :element-type (element-type m)))
	 (z0 (m* -1 (ones-vector (matrix-rows m) :element-type (element-type m)))))
    (augment-matrix w (mneg m) z0 q)))

(defun find-minimum-q (tableu)
  (let* ((jq (1- (matrix-cols tableu)))
	 (qmin (mref tableu 0 jq))
	 (qidx 0))
    (iter (for i from 1 below (matrix-rows tableu))
	  (finally (return qidx))
	  (if (< (mref tableu i jq) qmin)
	      (setf qmin (mref tableu i jq)
		    qidx i)))))

(defun find-minimum-ratio (tableu col)
  (let* ((first-nonzero (iter (for i from 0 below (matrix-rows tableu))
			      (finding i such-that (not (< (abs (mref tableu i col)) +epsilon+)))))
	 (jq (1- (matrix-cols tableu)))
	 (qmin (/ (mref tableu first-nonzero jq) (mref tableu first-nonzero col)))
	 (qidx first-nonzero))
    (iter (for i from qidx below (matrix-rows tableu))
	  (finally (return (if (minusp qmin)
			       (signal 'ray-termination)
			       (values qidx qmin))))
	  (unless (< (abs (mref tableu i col)) +epsilon+)
	    (let ((tst (/ (mref tableu i jq) (mref tableu i col))))
	      (unless (minusp tst)
		(if (< tst qmin)
		    (setf qmin tst
			  qidx i))))))))

(defun basis-vector-p (vec)
  (let ((0cnt 0)
	(1cnt 0)
	(len 0))
    (iter (for x in-vector vec)
	  (incf len)
	  (when (< (abs x) +epsilon+) (incf 0cnt))
	  (when (< (abs (- x 1.0d0)) +epsilon+) (incf 1cnt)))
    (and (= 1cnt 1) (= (+ 1cnt 0cnt) len))))

(defun find-q-basis-inverse (tableu)
  (let* ((vecs (append 
		(list (get-column tableu (1- (matrix-cols tableu))))
		(iter (for j from 0 below (1- (/ (- (matrix-cols tableu) 2) 2)))
		      (collect (get-column tableu j)))
		(list (get-column tableu (- (matrix-cols tableu) 2)))))
	 (basis (make-matrix (matrix-rows tableu) (length vecs)
			     :element-type (element-type tableu))))
    (iter (for c in vecs)
	  (for j from 0)
	  (set-column basis j c)
	  (finally (return basis)))))

(defun lexico-minimum-ratio-test (tableu entering)
  (let ((q-basis (find-q-basis-inverse tableu)))
    (lexico-minimum-idx (iter (for r in-matrix-row q-basis)
			      (for i from 0)
			      (let ((ais (mref tableu i entering)))
				(if (minusp ais)
				    (collect (m* most-positive-double-float (ones-vector (matrix-cols q-basis) :element-type (element-type tableu))))
				    (collect (m/ r (mref tableu i entering)))))))))

(defun find-dropping-variable (tableu entering &key (use-lexico-ratio-test nil))
  (if use-lexico-ratio-test
      (lexico-minimum-ratio-test tableu entering)
      (find-minimum-ratio tableu entering)))

(defun update-basic-variables (dict entering leaving)
  (prog1
      (svref dict leaving)
    (setf (svref dict leaving) entering)))

(defun nonpositive-p (v)
  (iter (for x in-vector v)
	(always (or (minusp x) (< (abs x) +epsilon+)))))

(defun lemke-step (tableu basic-vars leaving-var z1-offset use-lexico-minimum-ratio-method)
  (let* ((entering-col (if (< leaving-var z1-offset)
			   (+ z1-offset leaving-var)
			   (- leaving-var z1-offset)))
	 (dropping ;(find-minimum-ratio tableu entering-col)))
	  (find-dropping-variable tableu entering-col :use-lexico-ratio-test use-lexico-minimum-ratio-method)))
    (gaussian-pivot tableu dropping entering-col)
    (update-basic-variables basic-vars entering-col dropping)))

(defun lemke-method (m q &optional (use-lexico-minimum-ratio-method nil))
  "Solve a Linear Complimentary Problem using Lemke's linear complimentary pivot method."
  (let* ((tableu (build-tableu (copy-matrix m) q))
	 (z1-offset (matrix-cols m))
	 (z0 (- (matrix-cols tableu) 2))
	 (basic-vars (make-array (matrix-cols m) :initial-element 0))
	 (current-leaving 0)
	 (imin (find-minimum-q tableu)))
    (dotimes (i (matrix-cols m))
      (setf (svref basic-vars i) i))
    ;; Initialize
    (gaussian-pivot tableu imin z0)
    (setq current-leaving (update-basic-variables basic-vars z0 imin))
    ;; The pivot steps
    (iter (setq current-leaving (lemke-step tableu basic-vars current-leaving z1-offset use-lexico-minimum-ratio-method))
	  (finally (return (values basic-vars tableu)))
	  (until (= current-leaving z0)))))

(defun lcp-solver (m q)
  "Linear Complimentary Programming solver for problems formulated per Game Physics by Eberly, 2004."
  (multiple-value-bind (vars solution)
      (handler-case (lemke-method m q nil) ; First try using minimum ratio to determine pivot.
	(ray-termination () (lemke-method m q t))) ; Retry using lexico minimum ratio test to determine pivot.
    (flet ((one-position (v)
	     (iter (for x in-vector v)
		   (for i from 0)
		   (when (< (- 1.d0 x) +epsilon+)
		     (return i)))))
      (let ((q0 (1- (matrix-cols solution)))
	    (ret (zero-vector (* 2 (matrix-rows solution)))))
      (iter (for var in-vector vars)
	    (finally (return ret))
	    (let ((p (one-position (get-column solution var))))
	      (setf (aref ret var) (mref solution p q0))))))))
  
(defun solve-lp-as-lcp (a b c)
  (let* ((nw (matrix-rows a))
	 (nz (matrix-cols a))
	 (mtop (augment-matrix (zero-matrix nz nz) (transpose (copy-matrix a))))
	 (mbot (augment-matrix (mneg (copy-matrix a)) (zero-matrix nw nw)))
	 (m (make-matrix (+ nw nz) (+ nz nw)))
	 (q (zero-vector (+ (element-count b) (element-count c)))))
    (iter (for r in-matrix-row mtop)
	  (for i from 0)
	  (set-row m i r))
    (iter (for r in-matrix-row mbot)
	  (for i from nz)
	  (set-row m i r))
    (iter (for x in-vector c)
	  (for i from 0)
	  (setf (aref q i) (- x)))
    (iter (for x in-vector b)
	  (for i from nz)
	  (setf (aref q i) x))
    (lcp-solver m q)))
	 
    
;;; --------------------------------------

(defun test-tableu ()
    (%matrix '((1.0d0 1 -1 1 0 0 5)
	       (2 -3 1 0 1 0 3)
	       (-1 2 -1 0 0 1 1))))

(defun test-simplex-pivot ()
  (let ((tab (test-tableu)))
    (print-matrix (setf tab (simplex-pivot tab 0 0)))
    (format t "-----------------------~%")
    (print-matrix (setf tab (simplex-pivot tab 1 1)))
    (format t "-----------------------~%")
    (print-matrix (simplex-pivot tab 2 2))))

(defun eberly-eg-7.1 ()
  (%matrix '((-1.0d0  1 1 0 0  2)
	     ( 2 -1 0 1 0 -1)
	     ( 3  1 0 0 1  3))))

(defun test-leaving ()
  (let ((a1 (%matrix '((-3.0d0 2))))
	(a2 (%matrix '((6.0d0 6)))))
    (find-leaving-variable a2 a1)))

;; The moment of Truth
(defun test-simplex-solver ()
  (let ((a (%matrix '((1.0d0 2)
		      (2 1))))
	(c (%matrix '((1.0d0 1))))
	(b (%matrix '((6.0d0)
		      (6)))))
    (naive-simplex-solver a c b 2 2)))

;; WTF?  Eberly comes up with an xb of [1/4 9/4 3/4].  I get [1/4 9/4
;; -3/4], which shouldn't be right since the constraints are sn>=0,
;; xn>=0.  The last value is a slack variable, so perhaps it should be
;; read as - (-3/4) since the augmented matrix is [A | -I]?
(defun test-eberly-eg-7.1 ()
  (let ((a (%matrix '((-1.0d0 1)
		      (2 -1)
		      (3 1))))
	(c (%matrix '((1.0d0 1))))
	(b (%matrix '((2.0d0)
		      (-1)
		      (3)))))
    (naive-simplex-solver a c b 2 3)))

(defun test-mlp-example ()
  (let ((a (%matrix '((1.0d0 5)
		      (5 -1)
		      (-1 -1))))
	(c (%matrix '((1.0d0 10))))
	(b (%matrix '((-15.0d0)
		      (-11)
		      (4)))))
    (naive-simplex-solver a c b 2 3)))

(defun test-lcp-tableu ()
  (build-tableu (mneg (%matrix '((-1.0d0 1 1 1)
				 (1 -1 1 1)
				 (-1 -1 -2 0)
				 (-1 -1 0 -2))))
		(make-vector 3.0d0 5 -9 -5)))

(defun test-lemke-step-init ()
  (setq tableu (test-lcp-tableu))
  (setq z0 (- (matrix-cols tableu) 2))
  (setq z1-offset 4)
  (setq basic-vars (make-array (matrix-cols m) :initial-element 0))
  (dotimes (i 4)
    (setf (svref basic-vars i) i))
  (let ((imin (find-minimum-q tableu)))
    (gaussian-pivot tableu imin z0)
    (setq current-leaving (update-basic-variables basic-vars z0 imin))))

(defun test-lcp-1 ()
  (let ((m (mneg (%matrix '((-1.0d0 1 1 1)
			    (1 -1 1 1)
			    (-1 -1 -2 0)
			    (-1 -1 0 -2)))))
	(q (make-vector 3.0d0 5 -9 -5)))
    (lcp-solver m q)))

(defun test-lcp-2 ()
  (let ((m (mneg (%matrix '((1.0d0 0 3)
			    (-1 2 5)
			    (2 1 2)))))
	(q (make-vector -3.0d0 -2 -1)))
    (lcp-solver m q)))

(defun test-lcp-3 ()
  (let ((m (%matrix '((1.0d0 2 0)
		      (0 1 2)
		      (2 0 1))))
	(q (make-vector -1.0d0 -1 -1)))
    (lcp-solver m q)))

(defun test-lcp-4-tableu ()
  (let ((m (%matrix '((1.0d0 0 0)
		      (2 1 0)
		      (2 2 1))))
	(q (make-vector -8.0d0 -12 -14)))
    (build-tableu m q)))

(defun test-lcp-4 ()
  (let ((m (%matrix '((1.0d0 0 0)
		      (2 1 0)
		      (2 2 1))))
	(q (make-vector -8.0d0 -12.0d0 -14.0d0)))
    (lcp-solver m q)))